package cn.heycm.tenant.common.datasource;

import cn.heycm.tenant.common.context.TenantContextHolder;
import cn.heycm.tenant.common.datasource.entity.DataSourceItem;
import cn.heycm.tenant.common.datasource.entity.DataSourceItemPool;
import cn.heycm.tenant.common.datasource.util.DataSourceUtil;
import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
import org.springframework.util.StringUtils;

import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * 多租户数据源
 *
 * @author hey
 * @version 1.0
 * @since 2024/8/24 22:54
 */
public class TenantDataSource extends AbstractRoutingDataSource {

    public TenantDataSource() {
        init();
    }

    /**
     * 获取当前租户
     *
     * @return 租户ID
     */
    @Override
    protected Object determineCurrentLookupKey() {
        return TenantContextHolder.get();
    }

    private void init() {

        // 设置租户ID对应数据源
        Map<Object, Object> target = new HashMap<>();

        List<DataSourceItem> list = getDataSourceItems();
        String defaultTenant = null;
        for (DataSourceItem item : list) {
            target.put(item.getTenant(), DataSourceUtil.createDataSource(item));
            if (item.isPrimary()) {
                defaultTenant = item.getTenant();
            }
        }

        if (!StringUtils.hasText(defaultTenant)) {
            throw new RuntimeException("Default datasource not configured.");
        }

        // 设置租户数据源
        this.setTargetDataSources(target);
        // 设置默认租户
        this.setDefaultTargetDataSource(target.get(defaultTenant));
    }

    private static List<DataSourceItem> getDataSourceItems() {
        // 从配置文件、缓存、Nacos中获取数据源配置

        DataSourceItem item1 = new DataSourceItem();
        item1.setTenant("tenant1");
        item1.setPrimary(true);
        item1.setHost("localhost");
        item1.setPort(3306);
        item1.setSchema("tenant1");
        item1.setUsername("root");
        item1.setPassword("QH#heycm.2024");
        item1.setItemPool(new DataSourceItemPool());
        item1.getItemPool().setPoolName("tenant1-hk-pool");

        DataSourceItem item2 = new DataSourceItem();
        item2.setTenant("tenant2");
        item2.setHost("localhost");
        item2.setPort(3306);
        item2.setSchema("tenant2");
        item2.setUsername("root");
        item2.setPassword("QH#heycm.2024");
        item2.setItemPool(new DataSourceItemPool());
        item2.getItemPool().setPoolName("tenant2-hk-pool");

        return List.of(item1, item2);
    }
}
